from motors.feetech import FeetechMotorsBus
from motors.configs import FeetechMotorsBusConfig
from motors.feetech import TorqueMode
from vr_monitor import VRMonitor
import time
import threading
import asyncio
import numpy as np
import math
import pygame

# 配置串口参数和电机信息
config1 = FeetechMotorsBusConfig(
    port="/dev/ttyACM0",  # 左机械臂串口
    motors={
        "base_rotation": [11, "sts3215"],     # 底盘旋转
        "shoulder": [12, "sts3215"],          # 肩部关节
        "elbow": [13, "sts3215"],            # 肘部关节
        "wrist_pitch": [14, "sts3215"],      # 腕部俯仰
        "wrist_roll": [15, "sts3215"],       # 腕部旋转
        "gripper": [16, "sts3215"],          # 夹爪
    }
)

config2 = FeetechMotorsBusConfig(
    port="/dev/ttyACM0",  # 右机械臂串口
    motors={
        "base_rotation": [1, "sts3215"],     # 底盘旋转
        "shoulder": [2, "sts3215"],          # 肩部关节
        "elbow": [3, "sts3215"],            # 肘部关节  
        "wrist_pitch": [4, "sts3215"],      # 腕部俯仰
        "wrist_roll": [5, "sts3215"],       # 腕部旋转
        "gripper": [6, "sts3215"],          # 夹爪
    }
)

class RobotArm:
    """6自由度机械臂控制类"""
    
    def __init__(self, config, name):
        self.name = name
        self.motors_bus = FeetechMotorsBus(config)
        
        # 舵机转换常数
        self.SERVO_CENTER = 2048  # 舵机中位 (0度)
        self.SERVO_SCALE = 1000.0 / 90.0  # 舵机单位/度 (≈11.111)
        
        # 机械臂DH参数 (根据实际机械臂调整)
        self.link_lengths = {
            'base_height': 0.1,      # 底座高度
            'upper_arm': 0.15,       # 大臂长度  
            'forearm': 0.12,         # 小臂长度
            'wrist_length': 0.08     # 腕部到末端长度
        }
        
        # 关节角度限制 (舵机位置值，2048为中位0度)
        # 1048 ≈ -90度, 2048 = 0度, 3048 ≈ +90度
        self.joint_limits = {
            'base_rotation': (548, 3548),       # 底盘旋转 ±135度
            'shoulder': (1048, 3048),           # 肩部 -90度到+90度
            'elbow': (1048, 3048),              # 肘部 -90度到+90度  
            'wrist_pitch': (1048, 3048),        # 腕部俯仰 -90度到+90度
            'wrist_roll': (548, 3548),          # 腕部旋转 ±135度
            'gripper': (2080, 2800)             # 夹爪开合范围
        }
        
        # 当前关节位置
        self.current_positions = {}
        self.target_positions = {}
        
        # 初始化目标位置为中位
        for joint in self.joint_limits:
            if joint == 'gripper':
                self.target_positions[joint] = 2080  # 夹爪稍微闭合
            else:
                self.target_positions[joint] = self.SERVO_CENTER  # 其他关节中位
    
    def angle_to_servo(self, angle_deg):
        """将角度(度)转换为舵机位置值"""
        return int(self.SERVO_CENTER + angle_deg * self.SERVO_SCALE)
    
    def servo_to_angle(self, servo_pos):
        """将舵机位置值转换为角度(度)"""
        return (servo_pos - self.SERVO_CENTER) / self.SERVO_SCALE
    
    def move_to_home_position(self):
        """移动到初始位置"""
        print(f"{self.name} 移动到初始位置...")
        for joint in self.target_positions:
            if joint == 'gripper':
                self.target_positions[joint] = 2080
            else:
                self.target_positions[joint] = self.SERVO_CENTER
        
        return self.write_positions()
    
    def connect(self):
        """连接机械臂"""
        self.motors_bus.connect()
        self.motors_bus.write("Torque_Enable", TorqueMode.ENABLED.value)
        print(f"{self.name} 机械臂已连接")
    
    def disconnect(self):
        """断开连接"""
        self.motors_bus.disconnect()
        print(f"{self.name} 机械臂已断开连接")
    
    def read_positions(self):
        """读取当前关节位置"""
        try:
            positions = self.motors_bus.read("Present_Position")
            joint_names = list(self.joint_limits.keys())
            for i, joint in enumerate(joint_names):
                if i < len(positions):
                    self.current_positions[joint] = positions[i]
            return self.current_positions
        except Exception as e:
            print(f"读取{self.name}位置失败: {e}")
            return self.current_positions
    
    def write_positions(self):
        """写入目标位置到机械臂"""
        try:
            # 限制关节角度在安全范围内
            clamped_positions = []
            for joint in self.joint_limits:
                pos = self.target_positions[joint]
                min_pos, max_pos = self.joint_limits[joint]
                clamped_pos = max(min_pos, min(max_pos, int(pos)))
                clamped_positions.append(clamped_pos)
            
            self.motors_bus.write("Goal_Position", clamped_positions)
            return True
        except Exception as e:
            print(f"写入{self.name}位置失败: {e}")
            return False
    
    def cartesian_to_joint_angles(self, x, y, z, roll=0, pitch=0):
        """
        笛卡尔坐标到关节角度的逆运动学求解
        x, y, z: 目标位置 (米)
        roll, pitch: 末端姿态角度 (度)
        返回: 关节角度字典
        """
        try:
            # 1. 底盘旋转角度 (base_rotation)
            base_angle = math.atan2(y, x)
            
            # 2. 计算在机械臂平面内的距离
            r_horizontal = math.sqrt(x*x + y*y)
            r_vertical = z - self.link_lengths['base_height']
            
            # 3. 考虑腕部长度，计算肩肘关节的目标点
            wrist_x = r_horizontal - self.link_lengths['wrist_length'] * math.cos(math.radians(pitch))
            wrist_z = r_vertical - self.link_lengths['wrist_length'] * math.sin(math.radians(pitch))
            
            # 4. 肩肘关节逆运动学 (2连杆问题)
            L1 = self.link_lengths['upper_arm']
            L2 = self.link_lengths['forearm']
            
            # 计算到腕部关节的距离
            distance = math.sqrt(wrist_x*wrist_x + wrist_z*wrist_z)
            # 检查是否在工作空间内
            if distance > (L1 + L2) or distance < abs(L1 - L2):
                print(f"目标点 ({x:.3f}, {y:.3f}, {z:.3f}) 超出工作空间")
                return None
            
            # 使用余弦定理计算肘关节角度
            cos_elbow = (L1*L1 + L2*L2 - distance*distance) / (2*L1*L2)
            cos_elbow = max(-1, min(1, cos_elbow))  # 限制在[-1,1]范围
            elbow_angle = math.pi - math.acos(cos_elbow)
            
            # 计算肩关节角度
            alpha = math.atan2(wrist_z, wrist_x)
            beta = math.atan2(L2*math.sin(elbow_angle), L1 + L2*math.cos(elbow_angle))
            shoulder_angle = alpha - beta
            
            # 5. 计算腕部俯仰角度
            wrist_pitch_angle = math.radians(pitch) - shoulder_angle - elbow_angle
            
            # 6. 腕部旋转直接使用输入的roll角度
            wrist_roll_angle = math.radians(roll)
            
            # 转换为舵机位置值
            positions = {
                'base_rotation': self.angle_to_servo(base_angle * 180/math.pi),
                'shoulder': self.angle_to_servo(shoulder_angle * 180/math.pi),
                'elbow': self.angle_to_servo(elbow_angle * 180/math.pi),
                'wrist_pitch': self.angle_to_servo(wrist_pitch_angle * 180/math.pi),
                'wrist_roll': self.angle_to_servo(wrist_roll_angle * 180/math.pi),
                'gripper': self.target_positions['gripper']  # 保持当前夹爪位置
            }
            
            return positions
            
        except Exception as e:
            print(f"逆运动学计算失败: {e}")
            return None

class VRRobotController:
    """VR机械臂控制器主类"""
    
    def __init__(self):
        # 初始化机械臂
        self.left_arm = RobotArm(config1, "左臂")
        self.right_arm = RobotArm(config2, "右臂")
        
        # 初始化VR监控器
        self.vr_monitor = VRMonitor()
        
        # 控制参数
        self.vr_scale = 0.5  # VR坐标到实际坐标的缩放
        self.position_filter_alpha = 0.1  # 位置滤波系数
        self.angle_filter_alpha = 0.2    # 角度滤波系数
        
        # ========== 新增：相对位置控制相关变量 ==========
        self.reference_set = False  # 是否已设置参考位置
        self.left_reference_pos = None  # 左手VR参考位置
        self.right_reference_pos = None  # 右手VR参考位置
        self.left_reference_angles = {'roll': 0, 'pitch': 0}  # 左手参考角度
        self.right_reference_angles = {'roll': 0, 'pitch': 0}  # 右手参考角度
        
        # 机械臂参考位置（当前位置作为参考）
        self.left_arm_reference_pos = np.array([0.2, -0.1, 0.3])
        self.right_arm_reference_pos = np.array([0.2, 0.1, 0.3])
        # =============================================
        
        # 滤波后的目标位置
        self.filtered_left_pos = self.left_arm_reference_pos.copy()
        self.filtered_right_pos = self.right_arm_reference_pos.copy()
        self.filtered_left_angles = {'roll': 0, 'pitch': 0}
        self.filtered_right_angles = {'roll': 0, 'pitch': 0}
        
        # 控制状态
        self.is_running = False
        self.left_active = False
        self.right_active = False
        
        # pygame初始化
        pygame.init()
        self.screen = pygame.display.set_mode((800, 600))
        pygame.display.set_caption("VR机械臂控制系统")
        font_path = "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc"
        self.font = pygame.font.Font(font_path, 16)
        # self.font = pygame.font.SysFont("Arial", 16)
        # self.font = pygame.font.SysFont('simHei',20)
        self.clock = pygame.time.Clock()
    
    def set_vr_reference(self):
        """设置VR手柄当前位置为参考位置"""
        dual_goals = self.vr_monitor.get_latest_goal_nowait()
        if not dual_goals:
            print("❌ 无法获取VR数据，无法设置参考位置")
            return False
        
        left_goal = dual_goals.get("left")
        right_goal = dual_goals.get("right")
        
        # 设置左手参考位置
        if left_goal and left_goal.target_position is not None:
            self.left_reference_pos = np.array(left_goal.target_position)
            if left_goal.wrist_roll_deg is not None:
                self.left_reference_angles['roll'] = left_goal.wrist_roll_deg
            if left_goal.wrist_flex_deg is not None:
                self.left_reference_angles['pitch'] = left_goal.wrist_flex_deg
            print(f"✅ 左手参考位置设置: {self.left_reference_pos}")
        
        # 设置右手参考位置
        if right_goal and right_goal.target_position is not None:
            self.right_reference_pos = np.array(right_goal.target_position)
            if right_goal.wrist_roll_deg is not None:
                self.right_reference_angles['roll'] = right_goal.wrist_roll_deg
            if right_goal.wrist_flex_deg is not None:
                self.right_reference_angles['pitch'] = right_goal.wrist_flex_deg
            print(f"✅ 右手参考位置设置: {self.right_reference_pos}")
        
        # 如果两个手柄都设置成功，标记为已设置
        if self.left_reference_pos is not None and self.right_reference_pos is not None:
            self.reference_set = True
            print("🎯 VR参考位置设置完成！现在可以开始相对控制")
            return True
        elif self.left_reference_pos is not None or self.right_reference_pos is not None:
            self.reference_set = True
            print("🎯 部分VR参考位置设置完成！现在可以开始相对控制")
            return True
        else:
            print("❌ 无法检测到VR手柄，请确保VR系统正常运行")
            return False
    
    def start(self):
        """启动控制系统"""
        print("启动VR机械臂控制系统...")
        
        # 连接机械臂
        self.left_arm.connect()
        self.right_arm.connect()
        
        # 移动到初始位置
        print("移动机械臂到初始位置...")
        self.left_arm.move_to_home_position()
        self.right_arm.move_to_home_position()
        time.sleep(2)
        
        # 启动VR监控
        vr_thread = threading.Thread(
            target=lambda: asyncio.run(self.vr_monitor.start_monitoring()), 
            daemon=True
        )
        vr_thread.start()
        time.sleep(2)  # 等待VR初始化
        
        self.is_running = True
        print("系统启动完成!")
        print("\n🎮 重要提示：")
        print("   请按 R 键设置VR手柄的当前位置为参考位置")
        print("   设置后才能开始相对位置控制")
        
        # 主控制循环
        self.main_loop()
    
    def stop(self):
        """停止控制系统"""
        self.is_running = False
        self.left_arm.disconnect()
        self.right_arm.disconnect()
        pygame.quit()
        print("系统已停止")
    
    def low_pass_filter(self, new_value, old_value, alpha):
        """低通滤波器"""
        return alpha * new_value + (1 - alpha) * old_value
    
    def process_vr_data(self):
        """处理VR数据 - 修改为相对位置控制"""
        if not self.reference_set:
            # 如果还没设置参考位置，不进行控制
            return
        
        dual_goals = self.vr_monitor.get_latest_goal_nowait()
        if not dual_goals:
            return
        
        left_goal = dual_goals.get("left")
        right_goal = dual_goals.get("right")
        
        # 处理左手控制器 -> 左机械臂 (相对位置控制)
        if left_goal and left_goal.target_position is not None and self.left_reference_pos is not None:
            vr_pos = np.array(left_goal.target_position)
            
            # 计算相对于参考位置的偏移
            vr_offset = vr_pos - self.left_reference_pos
            
            # VR坐标系转换到机械臂坐标系的偏移
            robot_offset = np.array([
                -vr_offset[2] * self.vr_scale,  # VR的z偏移 -> 机械臂的x偏移
                -vr_offset[0] * self.vr_scale,  # VR的x偏移 -> 机械臂的y偏移
                vr_offset[1] * self.vr_scale    # VR的y偏移 -> 机械臂的z偏移
            ])
            
            # 计算目标位置 = 参考位置 + 偏移
            target_pos = self.left_arm_reference_pos + robot_offset
            
            # 位置滤波
            self.filtered_left_pos = self.low_pass_filter(
                target_pos, self.filtered_left_pos, self.position_filter_alpha
            )
            
            # 姿态角度滤波 (相对角度)
            if left_goal.wrist_roll_deg is not None:
                roll_offset = left_goal.wrist_roll_deg - self.left_reference_angles['roll']
                target_roll = roll_offset
                self.filtered_left_angles['roll'] = self.low_pass_filter(
                    target_roll, 
                    self.filtered_left_angles['roll'], 
                    self.angle_filter_alpha
                )
            
            if left_goal.wrist_flex_deg is not None:
                pitch_offset = left_goal.wrist_flex_deg - self.left_reference_angles['pitch']
                target_pitch = pitch_offset
                self.filtered_left_angles['pitch'] = self.low_pass_filter(
                    target_pitch,
                    self.filtered_left_angles['pitch'],
                    self.angle_filter_alpha
                )
            
            # 计算逆运动学
            joint_positions = self.left_arm.cartesian_to_joint_angles(
                self.filtered_left_pos[0],
                self.filtered_left_pos[1], 
                self.filtered_left_pos[2],
                self.filtered_left_angles['roll'],
                self.filtered_left_angles['pitch']
            )
            
            if joint_positions:
                self.left_arm.target_positions.update(joint_positions)
                self.left_active = True
            
            # 夹爪控制
            if left_goal.metadata and 'trigger_active' in left_goal.metadata:
                if left_goal.metadata['trigger_active']:
                    self.left_arm.target_positions['gripper'] = 2800  # 张开
                else:
                    self.left_arm.target_positions['gripper'] = 2080  # 闭合
        else:
            self.left_active = False
        
        # 处理右手控制器 -> 右机械臂 (相对位置控制)
        if right_goal and right_goal.target_position is not None and self.right_reference_pos is not None:
            vr_pos = np.array(right_goal.target_position)
            
            # 计算相对于参考位置的偏移
            vr_offset = vr_pos - self.right_reference_pos
            
            # VR坐标系转换到机械臂坐标系的偏移
            robot_offset = np.array([
                -vr_offset[2] * self.vr_scale,  # VR的z偏移 -> 机械臂的x偏移
                -vr_offset[0] * self.vr_scale,  # VR的x偏移 -> 机械臂的y偏移
                vr_offset[1] * self.vr_scale    # VR的y偏移 -> 机械臂的z偏移
            ])
            
            # 计算目标位置 = 参考位置 + 偏移
            target_pos = self.right_arm_reference_pos + robot_offset
            
            # 位置滤波
            self.filtered_right_pos = self.low_pass_filter(
                target_pos, self.filtered_right_pos, self.position_filter_alpha
            )
            
            # 姿态角度滤波 (相对角度)
            if right_goal.wrist_roll_deg is not None:
                roll_offset = right_goal.wrist_roll_deg - self.right_reference_angles['roll']
                target_roll = roll_offset
                self.filtered_right_angles['roll'] = self.low_pass_filter(
                    target_roll,
                    self.filtered_right_angles['roll'],
                    self.angle_filter_alpha
                )
            
            if right_goal.wrist_flex_deg is not None:
                pitch_offset = right_goal.wrist_flex_deg - self.right_reference_angles['pitch']
                target_pitch = pitch_offset
                self.filtered_right_angles['pitch'] = self.low_pass_filter(
                    target_pitch,
                    self.filtered_right_angles['pitch'],
                    self.angle_filter_alpha
                )
            
            # 计算逆运动学
            joint_positions = self.right_arm.cartesian_to_joint_angles(
                self.filtered_right_pos[0],
                self.filtered_right_pos[1],
                self.filtered_right_pos[2], 
                self.filtered_right_angles['roll'],
                self.filtered_right_angles['pitch']
            )
            
            if joint_positions:
                self.right_arm.target_positions.update(joint_positions)
                self.right_active = True
            
            # 夹爪控制
            if right_goal.metadata and 'trigger_active' in right_goal.metadata:
                if right_goal.metadata['trigger_active']:
                    self.right_arm.target_positions['gripper'] = 2800  # 张开
                else:
                    self.right_arm.target_positions['gripper'] = 2080  # 闭合
        else:
            self.right_active = False
    
    def update_robots(self):
        """更新机械臂位置"""
        # 读取当前位置
        self.left_arm.read_positions()
        self.right_arm.read_positions()
        
        # 发送目标位置
        if self.left_active:
            success = self.left_arm.write_positions()
            if not success:
                print("左臂控制失败")
        
        if self.right_active:
            success = self.right_arm.write_positions()
            if not success:
                print("右臂控制失败")
    
    def draw_ui(self):
        """绘制用户界面"""
        self.screen.fill((0, 0, 0))
        
        y_offset = 10
        
        # 标题
        title = self.font.render("VR机械臂控制系统 - 相对位置控制", True, (255, 255, 255))
        self.screen.blit(title, (10, y_offset))
        y_offset += 40
        
        # VR连接状态
        dual_goals = self.vr_monitor.get_latest_goal_nowait()
        left_connected = dual_goals and dual_goals.get("has_left", False)
        right_connected = dual_goals and dual_goals.get("has_right", False)
        
        vr_status = self.font.render(f"VR状态: 左手{'✓' if left_connected else '✗'} 右手{'✓' if right_connected else '✗'}", True, (0, 255, 0) if (left_connected or right_connected) else (255, 0, 0))
        self.screen.blit(vr_status, (10, y_offset))
        y_offset += 30
        
        # ========== 新增：参考位置设置状态显示 ==========
        ref_status_color = (0, 255, 0) if self.reference_set else (255, 165, 0)
        ref_status_text = "参考位置: ✓ 已设置" if self.reference_set else "参考位置: ⚠ 未设置 (按R键设置)"
        ref_status = self.font.render(ref_status_text, True, ref_status_color)
        self.screen.blit(ref_status, (10, y_offset))
        y_offset += 35
        # ============================================
        
        # 机械臂状态
        left_status = self.font.render(f"左臂: {'激活' if self.left_active else '待机'}", True, (0, 255, 0) if self.left_active else (128, 128, 128))
        self.screen.blit(left_status, (10, y_offset))
        y_offset += 25
        
        right_status = self.font.render(f"右臂: {'激活' if self.right_active else '待机'}", True, (0, 255, 0) if self.right_active else (128, 128, 128))
        self.screen.blit(right_status, (10, y_offset))
        y_offset += 35
        
        # 目标位置显示
        left_pos_text = self.font.render(f"左臂目标: [{self.filtered_left_pos[0]:.3f}, {self.filtered_left_pos[1]:.3f}, {self.filtered_left_pos[2]:.3f}]", True, (255, 255, 255))
        self.screen.blit(left_pos_text, (10, y_offset))
        y_offset += 25
        
        right_pos_text = self.font.render(f"右臂目标: [{self.filtered_right_pos[0]:.3f}, {self.filtered_right_pos[1]:.3f}, {self.filtered_right_pos[2]:.3f}]", True, (255, 255, 255))
        self.screen.blit(right_pos_text, (10, y_offset))
        y_offset += 35
        
        # 关节位置显示 (显示舵机位置值和对应角度)
        if self.left_arm.current_positions:
            y_offset += 10
            left_title = self.font.render("左臂关节状态:", True, (200, 200, 255))
            self.screen.blit(left_title, (10, y_offset))
            y_offset += 20
            
            for joint, pos in self.left_arm.current_positions.items():
                if joint != 'gripper':
                    angle = self.left_arm.servo_to_angle(pos)
                    target_pos = self.left_arm.target_positions.get(joint, pos)
                    target_angle = self.left_arm.servo_to_angle(target_pos)
                    joint_text = f"{joint}: {pos:4d} ({angle:+6.1f}°) -> {target_pos:4d} ({target_angle:+6.1f}°)"
                else:
                    target_pos = self.left_arm.target_positions.get(joint, pos)
                    joint_text = f"{joint}: {pos:4d} -> {target_pos:4d}"
                
                joint_display = self.font.render(joint_text, True, (180, 180, 180))
                self.screen.blit(joint_display, (10, y_offset))
                y_offset += 18
        
        if self.right_arm.current_positions:
            y_offset += 10
            right_title = self.font.render("右臂关节状态:", True, (255, 200, 200))
            self.screen.blit(right_title, (10, y_offset))
            y_offset += 20
            
            for joint, pos in self.right_arm.current_positions.items():
                if joint != 'gripper':
                    angle = self.right_arm.servo_to_angle(pos)
                    target_pos = self.right_arm.target_positions.get(joint, pos)
                    target_angle = self.right_arm.servo_to_angle(target_pos)
                    joint_text = f"{joint}: {pos:4d} ({angle:+6.1f}°) -> {target_pos:4d} ({target_angle:+6.1f}°)"
                else:
                    target_pos = self.right_arm.target_positions.get(joint, pos)
                    joint_text = f"{joint}: {pos:4d} -> {target_pos:4d}"
                
                joint_display = self.font.render(joint_text, True, (180, 180, 180))
                self.screen.blit(joint_display, (10, y_offset))
                y_offset += 18
        
        # 操作提示
        y_offset += 10
        instructions = [
            "操作说明:",
            "• R键 - 设置VR手柄当前位置为参考位置 ★",
            "• 移动VR控制器进行相对位置控制",
            "• 扳机键控制夹爪开合", 
            "• H键 - 回到初始位置",
            "• ESC键 - 退出程序",
            "",
            "舵机位置说明:",
            "• 2048 = 0度 (中位)",
            "• 3048 ≈ +90度",
            "• 1048 ≈ -90度"
        ]
        
        for instruction in instructions:
            color = (255, 255, 0) if "R键" in instruction else (180, 180, 180)
            text = self.font.render(instruction, True, color)
            self.screen.blit(text, (10, y_offset))
            y_offset += 20
        
        pygame.display.flip()
    
    def main_loop(self):
        """主控制循环"""
        while self.is_running:
            # 处理pygame事件
            for event in pygame.event.get():
                if event.type == pygame.QUIT or (event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE):
                    self.is_running = False
                    break
                elif event.type == pygame.KEYDOWN and event.key == pygame.K_h:
                    # H键回到初始位置
                    print("回到初始位置...")
                    self.left_arm.move_to_home_position()
                    self.right_arm.move_to_home_position()
                    self.left_active = False
                    self.right_active = False
                elif event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                    # ========== 新增：R键设置参考位置 ==========
                    print("🎯 正在设置VR参考位置...")
                    success = self.set_vr_reference()
                    if success:
                        print("✅ 参考位置设置成功！现在可以开始相对控制")
                    else:
                        print("❌ 参考位置设置失败，请检查VR连接")
                    # ======================================
            
            try:
                # 处理VR数据
                self.process_vr_data()
                
                # 更新机械臂
                self.update_robots()
                
                # 绘制界面
                self.draw_ui()
                
                # 控制循环频率
                self.clock.tick(50)  # 50Hz控制频率
                
            except Exception as e:
                print(f"控制循环错误: {e}")
                continue
        
        self.stop()

def main():
    """主函数"""
    controller = VRRobotController()
    try:
        controller.start()
    except KeyboardInterrupt:
        print("\n用户中断程序")
    except Exception as e:
        print(f"程序错误: {e}")
    finally:
        controller.stop()

if __name__ == "__main__":
    main()